// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#include <pollux/functions/prestosql/types/bing_tile_registration.h>

#include <pollux/expression/cast_expr.h>
#include <pollux/functions/prestosql/types/bing_tile_type.h>

namespace kumo::pollux {
    namespace {
        class BingTileCastOperator : public exec::CastOperator {
        public:
            static const std::shared_ptr<const CastOperator> &get() {
                static const std::shared_ptr<const CastOperator> instance{
                    new BingTileCastOperator()
                };

                return instance;
            }

            bool isSupportedFromType(const TypePtr &other) const override {
                switch (other->kind()) {
                    case TypeKind::BIGINT:
                        return true;
                    default:
                        return false;
                }
            }

            bool isSupportedToType(const TypePtr &other) const override {
                switch (other->kind()) {
                    case TypeKind::BIGINT:
                        return true;
                    default:
                        return false;
                }
            }

            void castTo(
                const BaseVector &input,
                exec::EvalCtx &context,
                const SelectivityVector &rows,
                const TypePtr &resultType,
                VectorPtr &result) const override {
                context.ensure_writable(rows, resultType, result);
                if (input.type_kind() == TypeKind::BIGINT) {
                    auto *bingTileResult = result->as_flat_vector<int64_t>();
                    bingTileResult->clear_nulls(rows);
                    const auto inputVector = input.as_checked<SimpleVector<int64_t> >();
                    // The values we just copy, and we will set errors later
                    bingTileResult->copy(inputVector, rows, nullptr);

                    // These are simple ops on int64s. We can vectorize for performance.
                    // For example, we could:
                    // 1. check validity isBingTileIntValid in parallel (branchless SIMD)
                    // 2. bitwise-OR valid flags with `rows` selectivity (branchless SIMD)
                    // 3. add an error to any false rows (uncommon)
                    context.applyToSelectedNoThrow(rows, [&](auto row) {
                        const uint64_t tileInt = bingTileResult->value_at(row);
                        if (MELON_UNLIKELY(!BingTileType::isBingTileIntValid(tileInt))) {
                            std::optional<std::string> reasonOpt =
                                    BingTileType::bingTileInvalidReason(tileInt);
                            if (reasonOpt.has_value()) {
                                context.setStatus(row, Status::UserError(*reasonOpt));
                            } else {
                                // isBingTileIntValid and bingTileInvalidReason have gotten out of
                                // sync
                                context.setStatus(
                                    row,
                                    Status::UnknownError(
                                        "Unknown BingTile validation error for tile %ld: this is a bug in Pollux",
                                        tileInt));
                            }
                        }
                    });
                } else {
                    POLLUX_UNSUPPORTED(
                        "Cast from {} to BINGTILE not supported", input.type()->toString());
                }
            }

            void castFrom(
                const BaseVector &input,
                exec::EvalCtx &context,
                const SelectivityVector &rows,
                const TypePtr &resultType,
                VectorPtr &result) const override {
                context.ensure_writable(rows, resultType, result);

                if (resultType->kind() == TypeKind::BIGINT) {
                    auto *flatResult = result->as_checked<FlatVector<int64_t> >();
                    flatResult->clear_nulls(rows);
                    const auto inputVector = input.as_checked<SimpleVector<int64_t> >();
                    flatResult->copy(inputVector, rows, nullptr);
                } else {
                    POLLUX_UNSUPPORTED(
                        "Cast from BINGTILE to {} not yet supported", resultType->toString());
                }
            }

        private:
            BingTileCastOperator() = default;
        };

        class BingTileTypeFactories : public CustomTypeFactories {
        public:
            pollux::TypePtr getType(
                const std::vector<pollux::TypeParameter> &parameters) const override {
                POLLUX_CHECK(parameters.empty());
                return BINGTILE();
            }

            exec::CastOperatorPtr getCastOperator() const override {
                return BingTileCastOperator::get();
            }

            AbstractInputGeneratorPtr getInputGenerator(
                const InputGeneratorConfig & /*config*/) const override {
                return nullptr;
            }
        };
    } // namespace

    void registerBingTileType() {
        registerCustomType(
            "bingtile", std::make_unique<const BingTileTypeFactories>());
    }
} // namespace kumo::pollux
